import threading
import concurrent.futures
import sys
import os
import time
import logging
import random
import numpy as np

DEOCODER_BASE_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, DEOCODER_BASE_DIR)
from asr_stream import StreamingDecoder, worker_feat, worker_decode
from AsrStDecoder import warp_asr_alloc_stdecoder, warp_asr_free_stdecoder,warp_asr_input_online,warp_asr_input_end,warp_asr_reset,warp_asr_gethyps,warp_asr_gethyps_num,warp_asr_asyn_decode
g_acouexe = None
g_asrmodel = None
g_decodermng = None
g_decodermng_lock = None

DECODER_TIMEOUT = 30

MODEL_PATH=bytes(DEOCODER_BASE_DIR, 'utf-8') +b'/model_zh'

def set_log():
    logger = logging.getLogger()
    logger.setLevel(logging.INFO)
    formatter = logging.Formatter('%(asctime)s | (%(module)s:%(lineno)d)| %(levelname)s | %(message)s')
    stdout_handler = logging.StreamHandler(sys.stdout)
    stdout_handler.setLevel(logging.INFO)
    stdout_handler.setFormatter(formatter)

    file_handler = logging.FileHandler('logs.log')
    file_handler.setLevel(logging.WARNING)
    file_handler.setFormatter(formatter)

    logger.addHandler(file_handler)
    logger.addHandler(stdout_handler)

def get_char_list(dictpath):
    with open(dictpath,'r') as f:
        alldict = f.readlines()
    char_list = ['<blank>']
    for linedict in alldict:
        linedict = linedict.strip().split()
        char_list.append(linedict[0])
    return char_list

class ASR_MODEL:
    def __init__(self, modelpath, modelnum,lang):
        self.sos = 1
        self.eos = 2
        self.beam = 5
        self.lang = lang
        self.token_list = get_char_list(os.path.join(modelpath,b'asr.dict'))
        self.max_num = modelnum
        self.model_list = []
        for i in range(self.max_num):
            stDecoder = np.zeros((1), dtype=np.int64)
            warp_asr_alloc_stdecoder(stDecoder, modelpath, b'model2',bytes(lang,'utf-8'))
            self.model_list.append({'sid':None,'encoder':stDecoder})

        self._model_lock = threading.Lock()

    def alloc_model(self,sid):
        for i in range(self.max_num):
            if self.model_list[i]['sid'] is None:
                self.model_list[i]['sid'] = sid
                return self.model_list[i]['encoder']

        logging.warning("error: reach maximum concurrent session {}".format(self.max_num))
        return None

    def free_model(self,sid):
        for i in range(self.max_num):
            if self.model_list[i]['sid'] is not None and self.model_list[i]['sid'] == sid:
                self.model_list[i]['sid'] = None 
                warp_asr_reset(self.model_list[i]['encoder'])

    def get_tokens(self,tids):
        if self.lang=="zh":
            return "".join([self.token_list[hid] for hid in tids])
        else:
            return " ".join([self.token_list[hid] for hid in tids]).replace("@@ ","").replace("@@","")


def asr_initmodel(lang="zh", max_session=20):
    global g_acouexe
    global g_asrmodel
    global g_decodermng
    global g_decodermng_lock
    global MODEL_PATH
    set_log()

    if lang=='en':
        MODEL_PATH=bytes(DEOCODER_BASE_DIR, 'utf-8') +b'/model_enbin'

    g_acouexe = concurrent.futures.ThreadPoolExecutor(max_session * 2)
    g_asrmodel = ASR_MODEL(MODEL_PATH,max_session,lang)
    g_decodermng = dict()
    g_decodermng_lock = threading.Lock()

    logging.info('asr decoder init succeed!')

def asr_decode_timeout(id):
    while True:
        with g_decodermng_lock:
            if id not in g_decodermng:
                break
            else:
                decoder_mng = g_decodermng[id]
        if decoder_mng['silent'] == 1:
            asr_decode_finish(id)
            break
        else:
            decoder_mng['silent'] = 1
        time.sleep(DECODER_TIMEOUT)
def asr_catch_excep(future):
    logger = logging.getLogger(__name__)
    #if future.exception(timeout=None):
    #    print(future.exception(timeout=None))
    try:
        future.result()
    except Exception:
        logger.exception("Executor Exception")

def asr_decode_online(id, wavbytes, use_punc=False, stm_punc=False):
    global g_acouexe
    global g_asrmodel
    global g_decodermng
    global g_decodermng_lock
    global g_cmvn_path
    global g_ngram_wfst
 
    with g_decodermng_lock:
        if id not in g_decodermng:
            asr_model = g_asrmodel.alloc_model(id)
            if asr_model is None:
                return {'result':[], 'subword':[]}
            logging.info('new session {}'.format(id))
            decoder = StreamingDecoder(id, asr_model)
            dectimer = threading.Thread(target=asr_decode_timeout, args=(id,))
            dectimer.start()
            g_decodermng[id] = {'dec':decoder, 'timer':dectimer, 'silent':0}

            fu1 = g_acouexe.submit(worker_feat, g_decodermng[id]['dec'])
            fu1.add_done_callback(asr_catch_excep)
            fu2 = g_acouexe.submit(worker_decode, g_decodermng[id]['dec'])
            fu2.add_done_callback(asr_catch_excep)
            decoder_mng = g_decodermng[id]
        else:
            decoder_mng = g_decodermng[id]
            decoder_mng['silent'] = 0
    wavraw = np.frombuffer(wavbytes, dtype=np.int8)
    result = decoder_mng['dec'].decode_streaming_window(wavraw)
    if result is None:
        return {'result':[], 'subword':[]}
    else:
        rec_text = g_asrmodel.get_tokens(result)
        return {'result':[rec_text], 'subword':[]}

def asr_decode_finish(id, topk=1):
    global g_acouexe
    global g_asrmodel
    global g_decodermng
    global g_decodermng_lock
    global g_lang_eng

    with g_decodermng_lock:
        if id not in g_decodermng:
            return {'result':[], 'subword':[]}
        else:
            decoder_mng = g_decodermng[id]
    finishst = time.time()
    #logging.info("finish st {} {}".format(id, finishst))
    result = decoder_mng['dec'].decode_streaming_finish()
    logging.info("finish end {} {}".format(id, time.time()-finishst))
    if result is None:
        final_out = {'result':[], 'subword':[]}
    else:
        rec_text = g_asrmodel.get_tokens(result)
        final_out = {'result':[rec_text], 'subword':[]}

    with g_decodermng_lock:
        g_asrmodel.free_model(id)
        del g_decodermng[id]['dec']
        del g_decodermng[id]['timer']
        del g_decodermng[id]
    return final_out

